from dataclasses import dataclass
from typing import override

import numpy as np
from sklearn.cluster import KMeans

from .emotion import EmotionalClassifiedStage
from .interface import Processor

__all__ = ["ClusterStage", "KMeansClusterizationProcessor"]


@dataclass(frozen=True)
class ClusterStage(object):
    points: np.ndarray
    labels: np.ndarray
    centers: np.ndarray


class KMeansClusterizationProcessor(Processor[EmotionalClassifiedStage, ClusterStage]):
    """
    A processor that performs KMeans clusterization on the emotional classified stage.

    This processor uses the `KMeans` class from `sklearn` to perform clusterization on the emotional classified stage.
    """

    _config: dict

    def __init__(self, global_config: dict) -> None:
        self._config = global_config["clusterizer"]["kmeans"]

    @override
    def process(self, stage: EmotionalClassifiedStage) -> ClusterStage:
        X = np.array(stage.emotions)
        kmeans = KMeans(
            n_clusters=self._config["n_clusters"],
            random_state=self._config["random_state"],
            max_iter=self._config["max_iter"],
        )
        labels = kmeans.fit_predict(X)
        centers = kmeans.cluster_centers_
        return ClusterStage(X, labels, centers)
